package com.wfh.core.handler;

import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.wfh.ai.model.message.*;
import com.wfh.ai.tool.BaseTool;
import com.wfh.ai.tool.ToolManager;
import com.wfh.constant.AppConstant;
import com.wfh.core.builder.VueProjectBuilder;
import com.wfh.model.entity.User;
import com.wfh.model.enums.ChatHistoryMessageTypeEnum;
import com.wfh.service.ChatHistoryService;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

import java.io.File;
import java.util.HashSet;
import java.util.Objects;
import java.util.Set;

/**
 * @Title: JsonMessageStreamHandler
 * @Author wangfenghuan
 * @Package com.wfh.aigenerator.core.handler
 * @Date 2025/8/5 21:30
 * @description: json消息流处理器
 */
@Slf4j
@Component
public class JsonMessageStreamHandler {


    @Resource
    private VueProjectBuilder vueProjectBuilder;


    @Resource
    private ToolManager toolManager;

    /**
     * 处理流式响应
     * @param originFlux
     * @param chatHistoryService
     * @param appId
     * @param loginUser
     * @return
     */
    public Flux<String> handle(Flux<String> originFlux,
                              ChatHistoryService chatHistoryService,
                              Long appId, User loginUser){
        // 收集数据用于生成后端记忆格式
        StringBuilder chatHistoryBuilder = new StringBuilder();
        // 用于跟踪已经见过的工具id，判断是否是第一次调用
        Set<String> seenToolIds = new HashSet<>();
        return originFlux.mapNotNull(json -> {
            // 解析每一个json消息快
            return handleJsonMessageChunk(json, chatHistoryBuilder, seenToolIds);
        })
                .filter(StrUtil::isNotEmpty)
                .doOnComplete(() -> {
                    // 流式响应完成之后，添加ai的消息到对话历史
                    String aiRes = chatHistoryBuilder.toString();
                    chatHistoryService.addChatMessage(appId, aiRes, ChatHistoryMessageTypeEnum.AI.getValue(), loginUser.getId());
                    // 生成vue项目
                    String projectPath = AppConstant.CODE_OUTPUT_ROOT_DIR + File.separator + "vue_project_" + appId;
                    vueProjectBuilder.buildProjectAsync(projectPath, appId);
                })
                .doOnError(e -> {
                    // 如果ai回复失败，也要记录错误的消息
                    String errorMessage = "发生错误：" + e.getMessage();
                    chatHistoryService.addChatMessage(appId, errorMessage, ChatHistoryMessageTypeEnum.AI.getValue(), loginUser.getId());
                });
    }

    /**
     * 处理每一个json消息块
     * @param json
     * @param chatHistoryBuilder
     * @param seenToolIds
     * @return
     */
    private String handleJsonMessageChunk(String json, StringBuilder chatHistoryBuilder, Set<String> seenToolIds) {
        // 解析json
        StreamMessage streamMessage = JSONUtil.toBean(json, StreamMessage.class);
        StreamMessageTypeEnum typeEnum = StreamMessageTypeEnum.getEnumByValue(streamMessage.getType());
        switch (Objects.requireNonNull(typeEnum)) {
            // ai的回复内容
            case AI_RESPONSE -> {
                AiResponseMessage aiResponseMessage = JSONUtil.toBean(json, AiResponseMessage.class);
                String data = aiResponseMessage.getData();
                // 直接拼接响应
                chatHistoryBuilder.append(data);
                return data;
            }
            // 调用工具
            case TOOL_REQUEST -> {
                ToolRequestMessage toolRequestMessage = JSONUtil.toBean(json, ToolRequestMessage.class);
                String toolId = toolRequestMessage.getId();
                String toolName = toolRequestMessage.getName();
                // 检查是否是第一次调用这个工具，记录id并返回完整的工具信息
                if (toolId != null && !seenToolIds.contains(toolId)){
                    seenToolIds.add(toolId);
                    BaseTool tool = toolManager.getTool(toolName);
                    return tool.generateToolRequestResponse();
                }else {
                    // 不是第一次调用，直接返回空字符串
                    return "";
                }
            }
            // 工具执行结果
            case TOOL_EXECUTED -> {
                ToolExecutedMessage toolExecutedMessage = JSONUtil.toBean(json, ToolExecutedMessage.class);
                JSONObject jsonObject = JSONUtil.parseObj(toolExecutedMessage.getArguments());
                String toolName = toolExecutedMessage.getName();
                BaseTool tool = toolManager.getTool(toolName);
                String result = tool.generateToolExecutedResult(jsonObject);
                // 输出前端和要持久化的内容
                String output = String.format("\n\n%s\n\n", result);
                chatHistoryBuilder.append(output);
                return output;
            }
            default -> {
                log.warn("未知的消息类型: {}", typeEnum);
                return "";
            }
        }
    }

}
